import os
import sys
import cv2
# from osgeo import gdal, ogr
import logging
import numpy as np
from PIL import Image
from tqdm import tqdm
import os
import numpy as np
from build.darknet.x64.detect1 import Detect
from build.darknet.x64 import darknet
from build.darknet.x64.nms import NMS


def TifCrop(SrcImg,CropSize, RepetitionRate,detect):

    # dataset_img = readTif(TifPath)
    # width = dataset_img.RasterXSize
    # height = dataset_img.RasterYSize
    # proj = dataset_img.GetProjection()
    # geotrans = dataset_img.GetGeoTransform()
    # img = dataset_img.ReadAsArray(0, 0, width, height)  # 获取数据
    #
    # print("++++++++++",img.shape)

    # img = cv2.imread(SrcPath)
    img = SrcImg
    height, width = img.shape[:2]


    new_detections = []


    #  获取当前文件夹的文件个数len,并以len+1命名即将裁剪得到的图像
    new_name = 1
    #  裁剪图片,重复率为RepetitionRate

    for i in range(int((height - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        for j in range(int((width - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
            #  如果图像是单波段
            if (len(img.shape) == 2):
                cropped = img[
                          int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                          int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
            #  如果图像是多波段
            else:
                cropped = img[
                          int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                          int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize,:]
            #  写图像
            # writeTiff(cropped, geotrans, proj, SavePath + '\\' +str(new_name) + '.tif')
            #  文件名 + 1

            new_detections.extend(detect.image_detection(cropped,
                                                         detect.network,
                                                         detect.class_names,
                                                         thresh=0.9,
                                                         x_offset=int(i * CropSize * (1 - RepetitionRate)),
                                                         y_offset=int(j * CropSize * (1 - RepetitionRate)),
                                                         ))

            new_name = new_name + 1
    #  向前裁剪最后一列
    for i in range(int((height - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        if (len(img.shape) == 2):
            cropped = img[int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                      (width - CropSize): width]
        else:
            cropped = img[
                      int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                      (width - CropSize): width,:]
        #  写图像
        # writeTiff(cropped, geotrans, proj, SavePath + '\\' +str(new_name)  + '.tif')
        new_detections.extend(detect.image_detection(cropped,
                                                     detect.network,
                                                     detect.class_names,
                                                     thresh=0.9,
                                                     x_offset=int(i * CropSize * (1 - RepetitionRate)),
                                                     y_offset=(width - CropSize),
                                                     ))
        new_name = new_name + 1
    #  向前裁剪最后一行
    for j in range(int((width - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        if (len(img.shape) == 2):
            cropped = img[(height - CropSize): height,
                      int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
        else:
            cropped = img[
                      (height - CropSize): height,
                      int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize,:]
        # writeTiff(cropped, geotrans, proj, SavePath + '\\' +str(new_name)  + '.tif')
        new_detections.extend(detect.image_detection(cropped,
                                                     detect.network,
                                                     detect.class_names,
                                                     thresh=0.9,
                                                     x_offset=(height - CropSize),
                                                     y_offset=int(j * CropSize * (1 - RepetitionRate)),
                                                    ))
        #  文件名 + 1
        new_name = new_name + 1
    #  裁剪右下角
    if (len(img.shape) == 2):
        cropped = img[(height - CropSize): height,
                  (width - CropSize): width]
    else:
        cropped = img[
                  (height - CropSize): height,
                  (width - CropSize): width,:]

    new_detections.extend(detect.image_detection(cropped,
                                                 detect.network,
                                                 detect.class_names,
                                                 thresh=0.5,
                                                 x_offset=(height - CropSize),
                                                 y_offset=(width - CropSize),
                                                 ))
    # writeTiff(cropped, geotrans, proj, SavePath + '\\' + str(new_name) + '_' + name + '.tif')



    return new_detections

def draw_boxes(detections, image, colors):
    import cv2
    for box in detections:
        left, top, right, bottom, confidence, label = box
        cv2.rectangle(image, (int(left),int(top)), (int(right),int(bottom)), colors[label], 1)
        # cv2.rectangle(image, (int(left),int(top)), (int(right),int(bottom)), colors[label], 1))
        cv2.putText(image, "{} [{:.2f}]".format(label, float(confidence)),
                    (int(left), int(top) - 5), cv2.FONT_HERSHEY_SIMPLEX, 5,
                    colors[label], 2)
    return image

if __name__ == '__main__':
    # gpu 通过环境变量设置

    detect = Detect(metaPath=r'G:/darknet/build/darknet/x64/train_data/obj.data',
                    configPath=r'G:/darknet/build/darknet/x64/train_data/yolov4-obj.cfg',
                    weightPath=r'G:/darknet/build/darknet/x64/train_data/back_up/yolov4-obj_last.weights',
                    gpu_id=0)


    #读取单张图像
    image_path = r'G:/object_detection_data/air_plane_test/hgw.tif'
    image = cv2.imread(image_path, -1)
    detections = TifCrop(image,800, 0.15 , detect)
    # detections = np.array(detections)
    print('detections nms 之前：',len(detections))
    nms = NMS()
    detections = nms(detections,detect.class_names)
    print('detections nms 之后：', len(detections))
    draw_bbox_image = draw_boxes(detections, image, detect.class_colors)
    cv2.imwrite('G:/object_detection_data/air_plane_test/hgw_pred1.tif',draw_bbox_image)

    # 读取文件夹
    # image_root = r'G:/darknet/build/darknet/x64/train_data/1/'
    # save_root = r'G:/darknet/build/darknet/x64/train_data/2'
    # if not os.path.exists(save_root):
    #     os.makedirs(save_root)
    # for name in os.listdir(image_root):
    #     print(name)
    #     image = cv2.imread(os.path.join(image_root, name), -1)
    #     draw_bbox_image = detect.predict_image(image, save_path=os.path.join(save_root, name))